import argparse
import math
import os
import random
import subprocess
import sys
import time
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path

try:
    import comet_ml  # must be imported before torch (if installed)
except ImportError:
    comet_ml = None

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

import val as validate  # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo_final import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (
    LOGGER,
    TQDM_BAR_FORMAT,
    check_amp,
    check_dataset,
    check_file,
    check_git_info,
    check_git_status,
    check_img_size,
    check_requirements,
    check_suffix,
    check_yaml,
    colorstr,
    get_latest_run,
    increment_path,
    init_seeds,
    intersect_dicts,
    labels_to_class_weights,
    labels_to_image_weights,
    methods,
    one_cycle,
    print_args,
    print_mutation,
    strip_optimizer,
    yaml_save,
)
from utils.loggers import LOGGERS, Loggers
from utils.loggers.comet.comet_utils import check_comet_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (
    EarlyStopping,
    ModelEMA,
    de_parallel,
    select_device,
    smart_DDP,
    smart_optimizer,
    smart_resume,
    torch_distributed_zero_first,
)

LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv("RANK", -1))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
GIT_INFO = check_git_info()


from EWC_module.EWC import compute_ewc_loss
from SI_module.SI_module import SynapticIntelligence
from Lwf_module.Lwf import distill_loss
from POD.POD import pod_spatial_loss





class Manager():
    def __init__(self, model, device, init_model_pack , paraList = None):
        cfg, hyp, exclude = init_model_pack
        self.device = device
        self.model = model
        self.paraList = paraList 
        '''
        应该是个嵌套字典。
        比如:{
            'Lwf':{'Lwf_enable':True, ...},
            ...
        }
        '''
        self.nc_old = self.paraList['nc_old']
        

        if self.paraList['Lwf']['Lwf_enable'] or self.paraList['POD']['PODNet_enable']:
            self.old_models = []
            for path in self.paraList['Old_models']:
                ckpt_i = torch.load(path, map_location="cpu")
        
                nc_old = ckpt_i['model'].model[-1].nc
                old_model = Model(cfg or ckpt_i["model"].yaml, 
                                  ch=3, 
                                  nc=nc_old, 
                                  anchors=hyp.get("anchors")
                                 ).to(device)  # create
                csd_old = ckpt_i["model"].float().state_dict()
                csd_old = intersect_dicts(csd_old, old_model.state_dict(), exclude=exclude)  # intersect
                old_model.load_state_dict(csd_old, strict = False)  # load
                # 我们这里要求旧模型严格一样
                old_model.nc = nc_old 
                for param in old_model.parameters():
                    param.requires_grad_(False) 
                    
                self.old_models.append((old_model, nc_old))

        
        self.SI_loss = {
            'old' : False, # 没有旧的si, if SI_loss['old']，则需要把损失加入
            'new' : False, # 不更新新的si
        }
        if self.paraList['SI']['SI_enable']:
            self.new_si_module = SynapticIntelligence(model, device, lambda_=self.paraList['SI']['SI_lambda'])
            if self.paraList['SI']['SI_pt'] is None:
                self.SI_loss['old'] = False
                # 说明现在是没有SImodel的，这是旧数据集。旧数据集训练不需要加入SIloss。
                 # 声明一个新的。
                pass
            else:
                # 说明这是新数据集。需要加入SIloss。
                self.SI_loss['old'] = True # 有旧数据集
                old_si_module = SynapticIntelligence(model, device, lambda_=self.paraList['SI']['SI_lambda'])
                #old_si_module = load            
                si_ckpt = torch.load(self.paraList['SI']['SI_pt'])
                old_si_module.omega = si_ckpt['omega']
                old_si_module.old_params = si_ckpt['old_params']

                self.old_si_module = old_si_module
                
        self.si_loss = []


            
        if self.paraList['EWC']['ewc_pt'] is not None:
            ewc_temp = torch.load(self.paraList['EWC']['ewc_pt'])
            self.ewc_pack = {
                'ewc_enable' : True,
                'fisher_matrix' : ewc_temp['fisher_matrix'],
                'old_modelpara' : ewc_temp['old_modelpara'],
                'model' : None, 
                'ewc_lambda': paraList['EWC']['ewc_lambda'], 
                'device':device, 
            }            
            for param in self.ewc_pack['old_modelpara'].values():
                param.requires_grad_(False)
            for param in self.ewc_pack['fisher_matrix'].values():
                param.requires_grad_(False)
        else:
            self.ewc_pack = {
                'ewc_enable' : False,
                'fisher_matrix' : None,
                'old_modelpara' : None,
                'model' : None, 
                'ewc_lambda': None, 
                'device':device, 
            }

    
    def cal_loss(self, imgs, targets, compute_loss, WORLD_SIZE, RANK, freeze, quad, scaler, optimizer, compute_loss_aux): 
        device = self.device
        model = self.model
        if self.paraList['DER']['DER_enable']:
            pred, features, aux_out = model(imgs, get_features = True) #这pred就是本次的输出。
            #aux_out = model.aux_out
        elif self.paraList['POD']['PODNet_enable']:
            pred, features = model(imgs, get_features = False, POD_features_layers = self.paraList['POD']['Distillation_layers'])
        else:
            pred = model(imgs)  # forward
        
        loss, loss_items = compute_loss(pred, 
                                        targets.to(device),
                                        )  # loss scaled by batch_size
        d_losses = []
        if self.paraList['POD']['PODNet_enable']:
            for (old_model, nc_old) in self.old_models:
                with torch.no_grad():
                    old_pred, old_features = old_model(imgs, 
                                                       get_features = False, 
                                                       POD_features_layers = self.paraList['POD']['Distillation_layers'])
                # 开始蒸馏，进行spatial级别的蒸馏，也就是宽高层面的，
                d_loss = []
                for n_f, o_f in zip(features, old_features):
                    d_loss.append(pod_spatial_loss(n_f, o_f))
                d_losses.append(sum(d_loss))
                loss += self.paraList['POD']['POD_lambda'] * sum(d_loss)
    
        if self.paraList['DER']['DER_enable']:
            aux_targets = targets.clone()
            for i in range(aux_targets.shape[0]):
                #print(aux_targets[i, 1])
                aux_targets[i, 1] = max(aux_targets[i, 1] - model.known_classes + 1, 0)
            aux_loss, aux_loss_items = compute_loss_aux(aux_out, 
                                            aux_targets.to(device),
                                            )  
    
            loss = loss + aux_loss

        if RANK != -1:
            loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode
        if quad:
            loss *= 4.0

            
        if self.paraList['Lwf']['Lwf_enable']:
            #print(len(self.old_models), len(self.paraList['Lwf']['Lwf_lambda']))
            for (old_model, nc_old), Lwf_lambda_i in zip(self.old_models, self.paraList['Lwf']['Lwf_lambda']):
                with torch.no_grad():
                    old_pred = old_model(imgs)
                new_pred = [pp[..., :5 + nc_old] for pp in pred] # 取共有的部分
                d_loss = distill_loss(new_pred, old_pred, nc_old, temperature = self.paraList['Lwf']['Lwf_temperature'])
                #print(Lwf_lambda_i * d_loss)
                loss += Lwf_lambda_i * d_loss
                d_losses.append(d_loss)
            

        if self.ewc_pack['ewc_enable']:
            self.ewc_pack['model'] = model
            ewc_loss = compute_ewc_loss(self.ewc_pack, method = 'L2', freeze = freeze)
            #print(f"EWC loss at batch {ni}: {ewc_loss.item()}")
            #ewc_loss.backward()
            loss += ewc_loss
            
        if self.SI_loss['old']:
            self.si_loss.append(self.old_si_module.penalty_loss(self.paraList['SI']['SI_lambda'], freeze = freeze))
            loss += self.si_loss[-1];

            
        scaler.scale(loss).backward()
        if self.paraList['SI']['SI_enable']:
            self.new_si_module.update_omega()#lr=optimizer.param_groups[0]['lr'])


        # end here
        return loss, d_losses, self.si_loss, loss_items

    
    




















